import numpy as np
from data_preparation.trajectory import Trajectory


class TrajectorySet:

    def __init__(self):
        self.trajectory_list = []
        self.trajectory_number = 0
        self.lmax = 0

    # this function give trajectory number
    def get_trajectory_number(self):
        return self.trajectory_number

    # this function add new trajectory in trajectory list
    def add_trajectory(self, trajectory1):
        sample1 = Trajectory()
        if type(sample1) is not type(trajectory1):
            raise TypeError('Must add trajectory to set')
        trajectory_number_now = self.get_trajectory_number()
        trajectory1.trajectory_index = trajectory_number_now + 1  # 为新添加的轨迹分配索引，索引值为当前轨迹数量加 1
        self.trajectory_list.append(trajectory1)
        self.trajectory_number = len(self.trajectory_list)

    # this function give trajectory according to index
    def give_trajectory_by_index(self, index1) -> Trajectory:
        try:
            trajectory1 = self.trajectory_list[index1]
        except IndexError:
            print(index1)
            raise IndexError

        return trajectory1

    # this function calculate all point number in this trajectory set
    def get_trajectory_set_point_number(self) -> int:
        point_number = 0
        trajectory_number = self.get_trajectory_number()
        for trajectory_index in range(trajectory_number):
            trajectory1 = self.give_trajectory_by_index(trajectory_index)
            this_trajectory_point_number = trajectory1.get_single_trajectory_point_number()
            point_number = point_number + this_trajectory_point_number
        return point_number

    def calculate_child_data(self, child_path: np.ndarray, node_tag: int) -> int:
        result = 0
        trajectory_number = self.get_trajectory_number()
        for trajectory_index in range(trajectory_number):
            trajectory1 = self.give_trajectory_by_index(trajectory_index)
            target_level2_index_sequence = trajectory1.after_interp_level2_unique_index_sequence
            length = target_level2_index_sequence.size
            if length < 1:
                length = 1
            window_size = len(child_path)
            if window_size == 0:
                print("window size is 0", child_path, node_tag)
                raise ValueError("window size is not valid")
            if window_size > len(target_level2_index_sequence):
                continue
            windows = np.lib.stride_tricks.sliding_window_view(target_level2_index_sequence, window_size)
            count = np.sum(np.all(windows == child_path, axis=1))
            for i in range(count):
                result = result + 1 / length
        return result

    def calculate_end_data(self,path : list) -> int:
        result = 0
        trajectory_number = self.get_trajectory_number()
        for trajectory_index in range(trajectory_number):
            trajectory1 = self.give_trajectory_by_index(trajectory_index)
            index_sequence = trajectory1.after_interp_level2_unique_index_sequence
            path_length = len(path)
            index_sequence_length = len(index_sequence)
            if np.array_equal(index_sequence[-path_length:],path):
                result = result + 1 /index_sequence_length
        return result


    # this function gives discrete trajectory the sample trajectory(unrepeated cell index array and its frequency)
    def get_unique_trajectory(self):
        for trajectory1 in self.trajectory_list:
            trajectory1.give_unique_trajectory()


